/*
 * Copyright 2024-2024 the original author or authors.
 */

package com.alibaba.langengine.modelcontextprotocol.client;

import com.alibaba.langengine.modelcontextprotocol.spec.*;
import com.alibaba.langengine.modelcontextprotocol.util.Assert;

import java.time.Duration;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.function.Consumer;
import java.util.function.Function;

/**
 * Factory class for creating Model Context Protocol (MCP) clients. MCP is a protocol that
 * enables AI models to interact with external tools and resources through a standardized
 * interface.
 *
 * <p>
 * This class serves as the main entry point for establishing connections with MCP
 * servers, implementing the client-side of the MCP specification. The protocol follows a
 * client-server architecture where:
 * <ul>
 * <li>The client (this implementation) initiates connections and sends requests
 * <li>The server responds to requests and provides access to tools and resources
 * <li>Communication occurs through a transport layer (e.g., stdio, SSE) using JSON-RPC
 * 2.0
 * </ul>
 *
 * <p>
 * The class provides factory methods to create either:
 * <ul>
 * <li>{@link McpAsyncClient} for non-blocking operations with CompletableFuture responses
 * <li>{@link McpSyncClient} for blocking operations with direct responses
 * </ul>
 *
 * <p>
 * Example of creating a basic synchronous client: <pre>{@code
 * McpClient.sync(transport)
 *     .requestTimeout(Duration.ofSeconds(5))
 *     .build();
 * }</pre>
 *
 * Example of creating a basic asynchronous client: <pre>{@code
 * McpClient.async(transport)
 *     .requestTimeout(Duration.ofSeconds(5))
 *     .build();
 * }</pre>
 *
 * <p>
 * Example with advanced asynchronous configuration: <pre>{@code
 * McpClient.async(transport)
 *     .requestTimeout(Duration.ofSeconds(10))
 *     .capabilities(new ClientCapabilities(...))
 *     .clientInfo(new Implementation("My Client", "1.0.0"))
 *     .roots(new Root("file://workspace", "Workspace Files"))
 *     .toolsChangeConsumer(tools -> {
 *         System.out.println("Tools updated: " + tools);
 *         return CompletableFuture.completedFuture(null);
 *     })
 *     .resourcesChangeConsumer(resources -> {
 *         System.out.println("Resources updated: " + resources);
 *         return CompletableFuture.completedFuture(null);
 *     })
 *     .promptsChangeConsumer(prompts -> {
 *         System.out.println("Prompts updated: " + prompts);
 *         return CompletableFuture.completedFuture(null);
 *     })
 *     .loggingConsumer(message -> {
 *         System.out.println("Log message: " + message);
 *         return CompletableFuture.completedFuture(null);
 *     })
 *     .build();
 * }</pre>
 *
 * <p>
 * The client supports:
 * <ul>
 * <li>Tool discovery and invocation
 * <li>Resource access and management
 * <li>Prompt template handling
 * <li>Real-time updates through change consumers
 * <li>Custom sampling strategies
 * <li>Structured logging with severity levels
 * </ul>
 *
 * <p>
 * The client supports structured logging through the MCP logging utility:
 * <ul>
 * <li>Eight severity levels from DEBUG to EMERGENCY
 * <li>Optional logger name categorization
 * <li>Configurable logging consumers
 * <li>Server-controlled minimum log level
 * </ul>
 *
 * JDK 1.8 compatible version.
 *
 * @author Christian Tzolov
 * @author Dariusz Jędrzejczyk
 * @see McpAsyncClient
 * @see McpSyncClient
 * @see McpTransport
 */
public interface McpClient {

    /**
     * Start building a synchronous MCP client with the specified transport layer. The
     * synchronous MCP client provides blocking operations. Synchronous clients wait for
     * each operation to complete before returning, making them simpler to use but
     * potentially less performant for concurrent operations. The transport layer handles
     * the low-level communication between client and server using protocols like stdio or
     * Server-Sent Events (SSE).
     * @param transport The transport layer implementation for MCP communication. Common
     * implementations include {@code StdioClientTransport} for stdio-based communication
     * and {@code SseClientTransport} for SSE-based communication.
     * @return A new builder instance for configuring the client
     * @throws IllegalArgumentException if transport is null
     */
    static SyncSpec sync(McpClientTransport transport) {
        return new SyncSpec(transport);
    }

    /**
     * Start building an asynchronous MCP client with the specified transport layer. The
     * asynchronous MCP client provides non-blocking operations. Asynchronous clients
     * return CompletableFuture immediately, allowing for concurrent
     * operations and asynchronous programming patterns. The transport layer handles the
     * low-level communication between client and server using protocols like stdio or
     * Server-Sent Events (SSE).
     * @param transport The transport layer implementation for MCP communication. Common
     * implementations include {@code StdioClientTransport} for stdio-based communication
     * and {@code SseClientTransport} for SSE-based communication.
     * @return A new builder instance for configuring the client
     * @throws IllegalArgumentException if transport is null
     */
    static AsyncSpec async(McpClientTransport transport) {
        return new AsyncSpec(transport);
    }

    /**
     * Synchronous client specification. This class follows the builder pattern to provide
     * a fluent API for setting up clients with custom configurations.
     *
     * <p>
     * The builder supports configuration of:
     * <ul>
     * <li>Transport layer for client-server communication
     * <li>Request timeouts for operation boundaries
     * <li>Client capabilities for feature negotiation
     * <li>Client implementation details for version tracking
     * <li>Root URIs for resource access
     * <li>Change notification handlers for tools, resources, and prompts
     * <li>Custom message sampling logic
     * </ul>
     */
    class SyncSpec {

        private final McpClientTransport transport;

        private Duration requestTimeout = Duration.ofSeconds(20); // Default timeout

        private Duration initializationTimeout = Duration.ofSeconds(20);

        private ClientCapabilities capabilities;

        private Implementation clientInfo = new Implementation("Java SDK MCP Client", "1.0.0");

        private final Map<String, Root> roots = new HashMap<>();

        private final List<Consumer<List<Tool>>> toolsChangeConsumers = new ArrayList<>();

        private final List<Consumer<List<Resource>>> resourcesChangeConsumers = new ArrayList<>();

        private final List<Consumer<List<Prompt>>> promptsChangeConsumers = new ArrayList<>();

        private final List<Consumer<LoggingMessageNotification>> loggingConsumers = new ArrayList<>();

        private Function<CreateMessageRequest, CreateMessageResult> samplingHandler;

        private SyncSpec(McpClientTransport transport) {
            Assert.notNull(transport, "Transport must not be null");
            this.transport = transport;
        }

        /**
         * Sets the duration to wait for server responses before timing out requests. This
         * timeout applies to all requests made through the client, including tool calls,
         * resource access, and prompt operations.
         * @param requestTimeout The duration to wait before timing out requests. Must not
         * be null.
         * @return This builder instance for method chaining
         * @throws IllegalArgumentException if requestTimeout is null
         */
        public SyncSpec requestTimeout(Duration requestTimeout) {
            Assert.notNull(requestTimeout, "Request timeout must not be null");
            this.requestTimeout = requestTimeout;
            return this;
        }

        /**
         * @param initializationTimeout The duration to wait for the initialization
         * lifecycle step to complete.
         * @return This builder instance for method chaining
         * @throws IllegalArgumentException if initializationTimeout is null
         */
        public SyncSpec initializationTimeout(Duration initializationTimeout) {
            Assert.notNull(initializationTimeout, "Initialization timeout must not be null");
            this.initializationTimeout = initializationTimeout;
            return this;
        }

        /**
         * Sets the client capabilities that will be advertised to the server during
         * connection initialization. Capabilities define what features the client
         * supports, such as tool execution, resource access, and prompt handling.
         * @param capabilities The client capabilities configuration. Must not be null.
         * @return This builder instance for method chaining
         * @throws IllegalArgumentException if capabilities is null
         */
        public SyncSpec capabilities(ClientCapabilities capabilities) {
            Assert.notNull(capabilities, "Capabilities must not be null");
            this.capabilities = capabilities;
            return this;
        }

        /**
         * Sets the client implementation information that will be shared with the server
         * during connection initialization. This helps with version compatibility and
         * debugging.
         * @param clientInfo The client implementation details including name and version.
         * Must not be null.
         * @return This builder instance for method chaining
         * @throws IllegalArgumentException if clientInfo is null
         */
        public SyncSpec clientInfo(Implementation clientInfo) {
            Assert.notNull(clientInfo, "Client info must not be null");
            this.clientInfo = clientInfo;
            return this;
        }

        /**
         * Sets the root URIs that this client can access. Roots define the base URIs for
         * resources that the client can request from the server. For example, a root
         * might be "file://workspace" for accessing workspace files.
         * @param roots A list of root definitions. Must not be null.
         * @return This builder instance for method chaining
         * @throws IllegalArgumentException if roots is null
         */
        public SyncSpec roots(List<Root> roots) {
            Assert.notNull(roots, "Roots must not be null");
            for (Root root : roots) {
                this.roots.put(root.uri(), root);
            }
            return this;
        }

        /**
         * Sets the root URIs that this client can access, using a varargs parameter for
         * convenience. This is an alternative to {@link #roots(List)}.
         * @param roots An array of root definitions. Must not be null.
         * @return This builder instance for method chaining
         * @throws IllegalArgumentException if roots is null
         * @see #roots(List)
         */
        public SyncSpec roots(Root... roots) {
            Assert.notNull(roots, "Roots must not be null");
            for (Root root : roots) {
                this.roots.put(root.uri(), root);
            }
            return this;
        }

        /**
         * Sets a custom sampling handler for processing message creation requests. The
         * sampling handler can modify or validate messages before they are sent to the
         * server, enabling custom processing logic.
         * @param samplingHandler A function that processes message requests and returns
         * results. Must not be null.
         * @return This builder instance for method chaining
         * @throws IllegalArgumentException if samplingHandler is null
         */
        public SyncSpec sampling(Function<CreateMessageRequest, CreateMessageResult> samplingHandler) {
            Assert.notNull(samplingHandler, "Sampling handler must not be null");
            this.samplingHandler = samplingHandler;
            return this;
        }

        /**
         * Adds a consumer to be notified when the available tools change. This allows the
         * client to react to changes in the server's tool capabilities, such as tools
         * being added or removed.
         * @param toolsChangeConsumer A consumer that receives the updated list of
         * available tools. Must not be null.
         * @return This builder instance for method chaining
         * @throws IllegalArgumentException if toolsChangeConsumer is null
         */
        public SyncSpec toolsChangeConsumer(Consumer<List<Tool>> toolsChangeConsumer) {
            Assert.notNull(toolsChangeConsumer, "Tools change consumer must not be null");
            this.toolsChangeConsumers.add(toolsChangeConsumer);
            return this;
        }

        /**
         * Adds a consumer to be notified when the available resources change. This allows
         * the client to react to changes in the server's resource availability, such as
         * files being added or removed.
         * @param resourcesChangeConsumer A consumer that receives the updated list of
         * available resources. Must not be null.
         * @return This builder instance for method chaining
         * @throws IllegalArgumentException if resourcesChangeConsumer is null
         */
        public SyncSpec resourcesChangeConsumer(Consumer<List<Resource>> resourcesChangeConsumer) {
            Assert.notNull(resourcesChangeConsumer, "Resources change consumer must not be null");
            this.resourcesChangeConsumers.add(resourcesChangeConsumer);
            return this;
        }

        /**
         * Adds a consumer to be notified when the available prompts change. This allows
         * the client to react to changes in the server's prompt templates, such as new
         * templates being added or existing ones being modified.
         * @param promptsChangeConsumer A consumer that receives the updated list of
         * available prompts. Must not be null.
         * @return This builder instance for method chaining
         * @throws IllegalArgumentException if promptsChangeConsumer is null
         */
        public SyncSpec promptsChangeConsumer(Consumer<List<Prompt>> promptsChangeConsumer) {
            Assert.notNull(promptsChangeConsumer, "Prompts change consumer must not be null");
            this.promptsChangeConsumers.add(promptsChangeConsumer);
            return this;
        }

        /**
         * Adds a consumer to be notified when logging messages are received from the
         * server. This allows the client to react to log messages, such as warnings or
         * errors, that are sent by the server.
         * @param loggingConsumer A consumer that receives logging messages. Must not be
         * null.
         * @return This builder instance for method chaining
         */
        public SyncSpec loggingConsumer(Consumer<LoggingMessageNotification> loggingConsumer) {
            Assert.notNull(loggingConsumer, "Logging consumer must not be null");
            this.loggingConsumers.add(loggingConsumer);
            return this;
        }

        /**
         * Adds multiple consumers to be notified when logging messages are received from
         * the server. This allows the client to react to log messages, such as warnings
         * or errors, that are sent by the server.
         * @param loggingConsumers A list of consumers that receive logging messages. Must
         * not be null.
         * @return This builder instance for method chaining
         */
        public SyncSpec loggingConsumers(List<Consumer<LoggingMessageNotification>> loggingConsumers) {
            Assert.notNull(loggingConsumers, "Logging consumers must not be null");
            this.loggingConsumers.addAll(loggingConsumers);
            return this;
        }

        /**
         * Create an instance of {@link McpSyncClient} with the provided configurations or
         * sensible defaults.
         * @return a new instance of {@link McpSyncClient}.
         */
        public McpSyncClient build() {
            // Default capabilities if none provided
            if (capabilities == null) {
                capabilities = ClientCapabilities.builder()
                        .build();
            }

            return new McpSyncClient(
                    transport,
                    requestTimeout,
                    initializationTimeout,
                    capabilities,
                    clientInfo,
                    new ArrayList<>(roots.values()),
                    toolsChangeConsumers,
                    resourcesChangeConsumers,
                    promptsChangeConsumers,
                    loggingConsumers,
                    samplingHandler);
        }
    }

    /**
     * Asynchronous client specification. This class follows the builder pattern to
     * provide a fluent API for setting up clients with custom configurations.
     *
     * <p>
     * The builder supports configuration of:
     * <ul>
     * <li>Transport layer for client-server communication
     * <li>Request timeouts for operation boundaries
     * <li>Client capabilities for feature negotiation
     * <li>Client implementation details for version tracking
     * <li>Root URIs for resource access
     * <li>Change notification handlers for tools, resources, and prompts
     * <li>Custom message sampling logic
     * </ul>
     */
    class AsyncSpec {

        private final McpClientTransport transport;

        private Duration requestTimeout = Duration.ofSeconds(20); // Default timeout

        private Duration initializationTimeout = Duration.ofSeconds(20);

        private ClientCapabilities capabilities;

        private Implementation clientInfo = new Implementation("Java SDK MCP Client", "1.0.0");

        private final Map<String, Root> roots = new HashMap<>();

        private final List<Function<List<Tool>, CompletableFuture<Void>>> toolsChangeConsumers = new ArrayList<>();

        private final List<Function<List<Resource>, CompletableFuture<Void>>> resourcesChangeConsumers = new ArrayList<>();

        private final List<Function<List<Prompt>, CompletableFuture<Void>>> promptsChangeConsumers = new ArrayList<>();

        private final List<Function<LoggingMessageNotification, CompletableFuture<Void>>> loggingConsumers = new ArrayList<>();

        private Function<CreateMessageRequest, CompletableFuture<CreateMessageResult>> samplingHandler;

        private AsyncSpec(McpClientTransport transport) {
            Assert.notNull(transport, "Transport must not be null");
            this.transport = transport;
        }

        /**
         * Sets the duration to wait for server responses before timing out requests. This
         * timeout applies to all requests made through the client, including tool calls,
         * resource access, and prompt operations.
         * @param requestTimeout The duration to wait before timing out requests. Must not
         * be null.
         * @return This builder instance for method chaining
         * @throws IllegalArgumentException if requestTimeout is null
         */
        public AsyncSpec requestTimeout(Duration requestTimeout) {
            Assert.notNull(requestTimeout, "Request timeout must not be null");
            this.requestTimeout = requestTimeout;
            return this;
        }

        /**
         * @param initializationTimeout The duration to wait for the initialization
         * lifecycle step to complete.
         * @return This builder instance for method chaining
         * @throws IllegalArgumentException if initializationTimeout is null
         */
        public AsyncSpec initializationTimeout(Duration initializationTimeout) {
            Assert.notNull(initializationTimeout, "Initialization timeout must not be null");
            this.initializationTimeout = initializationTimeout;
            return this;
        }

        /**
         * Sets the client capabilities that will be advertised to the server during
         * connection initialization. Capabilities define what features the client
         * supports, such as tool execution, resource access, and prompt handling.
         * @param capabilities The client capabilities configuration. Must not be null.
         * @return This builder instance for method chaining
         * @throws IllegalArgumentException if capabilities is null
         */
        public AsyncSpec capabilities(ClientCapabilities capabilities) {
            Assert.notNull(capabilities, "Capabilities must not be null");
            this.capabilities = capabilities;
            return this;
        }

        /**
         * Sets the client implementation information that will be shared with the server
         * during connection initialization. This helps with version compatibility and
         * debugging.
         * @param clientInfo The client implementation details including name and version.
         * Must not be null.
         * @return This builder instance for method chaining
         * @throws IllegalArgumentException if clientInfo is null
         */
        public AsyncSpec clientInfo(Implementation clientInfo) {
            Assert.notNull(clientInfo, "Client info must not be null");
            this.clientInfo = clientInfo;
            return this;
        }

        /**
         * Sets the root URIs that this client can access. Roots define the base URIs for
         * resources that the client can request from the server. For example, a root
         * might be "file://workspace" for accessing workspace files.
         * @param roots A list of root definitions. Must not be null.
         * @return This builder instance for method chaining
         * @throws IllegalArgumentException if roots is null
         */
        public AsyncSpec roots(List<Root> roots) {
            Assert.notNull(roots, "Roots must not be null");
            for (Root root : roots) {
                this.roots.put(root.uri(), root);
            }
            return this;
        }

        /**
         * Sets the root URIs that this client can access, using a varargs parameter for
         * convenience. This is an alternative to {@link #roots(List)}.
         * @param roots An array of root definitions. Must not be null.
         * @return This builder instance for method chaining
         * @throws IllegalArgumentException if roots is null
         * @see #roots(List)
         */
        public AsyncSpec roots(Root... roots) {
            Assert.notNull(roots, "Roots must not be null");
            for (Root root : roots) {
                this.roots.put(root.uri(), root);
            }
            return this;
        }

        /**
         * Sets a custom sampling handler for processing message creation requests. The
         * sampling handler can modify or validate messages before they are sent to the
         * server, enabling custom processing logic.
         * @param samplingHandler A function that processes message requests and returns
         * results. Must not be null.
         * @return This builder instance for method chaining
         * @throws IllegalArgumentException if samplingHandler is null
         */
        public AsyncSpec sampling(Function<CreateMessageRequest, CompletableFuture<CreateMessageResult>> samplingHandler) {
            Assert.notNull(samplingHandler, "Sampling handler must not be null");
            this.samplingHandler = samplingHandler;
            return this;
        }

        /**
         * Adds a consumer to be notified when the available tools change. This allows the
         * client to react to changes in the server's tool capabilities, such as tools
         * being added or removed.
         * @param toolsChangeConsumer A consumer that receives the updated list of
         * available tools. Must not be null.
         * @return This builder instance for method chaining
         * @throws IllegalArgumentException if toolsChangeConsumer is null
         */
        public AsyncSpec toolsChangeConsumer(Function<List<Tool>, CompletableFuture<Void>> toolsChangeConsumer) {
            Assert.notNull(toolsChangeConsumer, "Tools change consumer must not be null");
            this.toolsChangeConsumers.add(toolsChangeConsumer);
            return this;
        }

        /**
         * Adds a consumer to be notified when the available resources change. This allows
         * the client to react to changes in the server's resource availability, such as
         * files being added or removed.
         * @param resourcesChangeConsumer A consumer that receives the updated list of
         * available resources. Must not be null.
         * @return This builder instance for method chaining
         * @throws IllegalArgumentException if resourcesChangeConsumer is null
         */
        public AsyncSpec resourcesChangeConsumer(Function<List<Resource>, CompletableFuture<Void>> resourcesChangeConsumer) {
            Assert.notNull(resourcesChangeConsumer, "Resources change consumer must not be null");
            this.resourcesChangeConsumers.add(resourcesChangeConsumer);
            return this;
        }

        /**
         * Adds a consumer to be notified when the available prompts change. This allows
         * the client to react to changes in the server's prompt templates, such as new
         * templates being added or existing ones being modified.
         * @param promptsChangeConsumer A consumer that receives the updated list of
         * available prompts. Must not be null.
         * @return This builder instance for method chaining
         * @throws IllegalArgumentException if promptsChangeConsumer is null
         */
        public AsyncSpec promptsChangeConsumer(Function<List<Prompt>, CompletableFuture<Void>> promptsChangeConsumer) {
            Assert.notNull(promptsChangeConsumer, "Prompts change consumer must not be null");
            this.promptsChangeConsumers.add(promptsChangeConsumer);
            return this;
        }

        /**
         * Adds a consumer to be notified when logging messages are received from the
         * server. This allows the client to react to log messages, such as warnings or
         * errors, that are sent by the server.
         * @param loggingConsumer A consumer that receives logging messages. Must not be
         * null.
         * @return This builder instance for method chaining
         */
        public AsyncSpec loggingConsumer(Function<LoggingMessageNotification, CompletableFuture<Void>> loggingConsumer) {
            Assert.notNull(loggingConsumer, "Logging consumer must not be null");
            this.loggingConsumers.add(loggingConsumer);
            return this;
        }

        /**
         * Adds multiple consumers to be notified when logging messages are received from
         * the server. This allows the client to react to log messages, such as warnings
         * or errors, that are sent by the server.
         * @param loggingConsumers A list of consumers that receive logging messages. Must
         * not be null.
         * @return This builder instance for method chaining
         */
        public AsyncSpec loggingConsumers(List<Function<LoggingMessageNotification, CompletableFuture<Void>>> loggingConsumers) {
            Assert.notNull(loggingConsumers, "Logging consumers must not be null");
            this.loggingConsumers.addAll(loggingConsumers);
            return this;
        }

        /**
         * Create an instance of {@link McpAsyncClient} with the provided configurations
         * or sensible defaults.
         * @return a new instance of {@link McpAsyncClient}.
         */
        public McpAsyncClient build() {
            // Default capabilities if none provided
            if (capabilities == null) {
                capabilities = ClientCapabilities.builder()
                        .build();
            }

            return new McpAsyncClient(
                    transport,
                    requestTimeout,
                    initializationTimeout,
                    capabilities,
                    clientInfo,
                    new ArrayList<>(roots.values()),
                    toolsChangeConsumers,
                    resourcesChangeConsumers,
                    promptsChangeConsumers,
                    loggingConsumers,
                    samplingHandler);
        }
    }
}
